package mr

import (
	"fmt"
	"hash/fnv"
	"io/ioutil"
	"log"
	"net/rpc"
	"os"
	"sort"
	"strings"
	"time"
)

//
// Map functions return a slice of KeyValue.
//
type KeyValue struct {
	Key   string
	Value string
}

// for sorting by key.
type ByKey []KeyValue

// for sorting by key.
func (a ByKey) Len() int           { return len(a) }
func (a ByKey) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByKey) Less(i, j int) bool { return a[i].Key < a[j].Key }

//
// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
//
func ihash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

//
// main/mrworker.go calls this function.
//
func Worker(mapf func(string, string) []KeyValue,
	reducef func(string, []string) string) {

	// Your worker implementation here.

	tempDir := "temp"
	if _, err := os.Stat(tempDir); os.IsNotExist(err) {
		os.Mkdir(tempDir, 0745)
	}

	for {
		task := CallRequestTask()
		// log.Printf("%+v\n", task)

		if task.Type == MAP_TASK {
			outFiles := make(map[string]*os.File)
			outFilenames := []string{}
			for _, filename := range task.InputFilenames {
				buf, err := ioutil.ReadFile(filename)
				if err != nil {
					log.Println(err)
				}

				cnt := string(buf)
				kvs := mapf(filename, cnt)
				for _, kv := range kvs {
					out := fmt.Sprintf("mr-%v-%v", task.MapId, ihash(kv.Key)%task.NReduce)
					f, exists := outFiles[out]
					if !exists {
						nf, err := ioutil.TempFile(tempDir, out+"*")
						if err != nil {
							log.Println(err)
							continue
						}

						outFiles[out] = nf
						f = nf

						outFilenames = append(outFilenames, out)
					}
					_, err := f.WriteString(kv.Key + " " + kv.Value + "\n")
					if err != nil {
						log.Println(err)
					}
				}

			}

			for filename, outf := range outFiles {
				outf.Close()

				err := os.Rename(outf.Name(), filename)
				if err != nil {
					log.Println(err)
				}
			}

			task.Finished = true
			finishArgs := FinishArgs{
				Task: task.Task,
			}
			finishArgs.OutputFilenames = outFilenames
			CallFinishTask(finishArgs)
		} else if task.Type == REDUCE_TASK {
			kvs := []KeyValue{}
			for _, filename := range task.InputFilenames {
				buf, err := ioutil.ReadFile(filename)
				if err != nil {
					log.Println(err)
				}

				cnt := string(buf)

				lines := strings.FieldsFunc(cnt, func(r rune) bool { return r == '\n' })
				for _, line := range lines {
					kv := strings.Fields(line)
					if len(kv) < 2 {
						log.Println(line)
						log.Println(kv)
					}
					kvs = append(kvs, KeyValue{
						Key:   kv[0],
						Value: kv[1],
					})
				}
			}

			sort.Sort(ByKey(kvs))

			outFilename := fmt.Sprintf("mr-out-%v", task.ReduceId)
			ofile, err := ioutil.TempFile(tempDir, outFilename+"*")
			if err != nil {
				log.Println(err)
			}

			i := 0
			for i < len(kvs) {
				j := i + 1
				for j < len(kvs) && kvs[j].Key == kvs[i].Key {
					j++
				}
				values := []string{}
				for k := i; k < j; k++ {
					values = append(values, kvs[k].Value)
				}
				output := reducef(kvs[i].Key, values)

				// this is the correct format for each line of Reduce output.
				fmt.Fprintf(ofile, "%v %v\n", kvs[i].Key, output)

				i = j
			}
			ofile.Close()
			err = os.Rename(ofile.Name(), outFilename)
			if err != nil {
				log.Println(err)
			}

			task.Finished = true
			finishArgs := FinishArgs{
				Task: task.Task,
			}
			finishArgs.OutputFilenames = []string{outFilename}
			CallFinishTask(finishArgs)
		} else {
			time.Sleep(time.Second)
		}
	}

	// uncomment to send the Example RPC to the coordinator.
	// CallExample()
}

//
// example function to show how to make an RPC call to the coordinator.
//
// the RPC argument and reply types are defined in rpc.go.
//
func CallExample() {

	// declare an argument structure.
	args := ExampleArgs{}

	// fill in the argument(s).
	args.X = 99

	// declare a reply structure.
	reply := ExampleReply{}

	// send the RPC request, wait for the reply.
	call("Coordinator.Example", &args, &reply)

	// reply.Y should be 100.
	fmt.Printf("reply.Y %v\n", reply.Y)
}

func CallRequestTask() TaskReply {
	args := TaskArgs{}

	reply := TaskReply{}

	call("Coordinator.RequestTask", &args, &reply)

	return reply
}

func CallFinishTask(args FinishArgs) {
	reply := FinishReply{}

	call("Coordinator.FinishTask", &args, &reply)
}

//
// send an RPC request to the coordinator, wait for the response.
// usually returns true.
// returns false if something goes wrong.
//
func call(rpcname string, args interface{}, reply interface{}) bool {
	// c, err := rpc.DialHTTP("tcp", "127.0.0.1"+":1234")
	sockname := coordinatorSock()
	c, err := rpc.DialHTTP("unix", sockname)
	if err != nil {
		log.Fatal("dialing:", err)
	}
	defer c.Close()

	err = c.Call(rpcname, args, reply)
	if err == nil {
		return true
	}

	fmt.Println(err)
	return false
}
